# BTControl.py
#
# Copyright (c) 2006 Michael Hobbs
#
# This file is part of iTorrent.
#
# iTorrent is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# iTorrent is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with iTorrent; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA


import threading
import os.path
import time

from BitTorrent.download import Multitorrent
from BitTorrent import GetTorrent
from BitTorrent.bencode import bdecode
from BitTorrent.ConvertedMetainfo import ConvertedMetainfo
from BitTorrent.prefs import Preferences
from BitTorrent.Storage import FilePool, Storage
from BitTorrent.StorageWrapper import StorageWrapper
from BitTorrent import BTFailure, BTShutdown, INFO, WARNING, ERROR, CRITICAL

import utils


CHECK_SECONDS = 1  # Frequency for RawServer to check for pending additions
LEVEL_FUNC = {INFO: 'info', WARNING: 'warning', ERROR: 'error', CRITICAL: 'critical'}


class TorrentInfo(object):
	def __init__(self, metainfo, config, handler, keepAlive, url, finishTimeout):
		super(TorrentInfo, self).__init__()
		self.metainfo = metainfo
		self.config = config
		self.handler = handler
		self.keepAlive = keepAlive
		self.url = url
		# TODO: look in config['save_in'] for destination path? (only if keep_downloads is set)
		if not os.path.isdir(utils.PARTIALS_DIR):
			os.makedirs(utils.PARTIALS_DIR)
		self.pathName = utils.getPartialPathName(metainfo.name_fs)
		self.finishTimeout = finishTimeout
		self.started = None


class BTControl(object):
	def __init__(self, config, display):
		super(BTControl, self).__init__()
		self.config = config
		self.display = display
		self.doneFlag = threading.Event()
		self.multitorrent = Multitorrent(self.config, self.doneFlag, self.error)
		self.rawServer = self.multitorrent.rawserver
		self.torrentInfo = {}
		self.sharingTorrents = {}
		# Add initial tasks to rawServer:
		self.updateDisplay()

	def start(self):
		self.rawServer.listen_forever()

	def stop(self):
		self.doneFlag.set()

	def getTorrent(self, data, url, handler=None, pieces=[], keepAlive=None, finishTimeout=0):
		metainfo = ConvertedMetainfo(bdecode(data))
		config = self.getConfig(pieces, keepAlive)
		info = TorrentInfo(metainfo, config, handler, keepAlive, url, finishTimeout)
		torrent = self.multitorrent.start_torrent(metainfo, config, self, info.pathName)
		config['#torrent'] = torrent
		self.torrentInfo[torrent] = info
		return (info.pathName, torrent)

	def shutdownTorrent(self, torrent):
		self.display.shutdownTorrent(torrent)
		torrent.shutdown()

	def stopTorrent(self, torrent):
		self.torrentInfo.pop(torrent, None)
		self.sharingTorrents.pop(torrent, None)
		# TODO: implement a stoppedTorrent method in display
		self.shutdownTorrent(torrent)

	def startKeepAlive(self, torrentInfo, torrent):
		if not torrentInfo.keepAlive:
			return
		def serverCallback():
			# Check if the torrent has stopped
			if torrent not in self.torrentInfo:
				return
			self.rawServer.add_task(serverCallback, 1)
			torrentInfo.keepAlive.keepAlive(torrent)
		serverCallback()

	# TODO: rename this to something more appropriate, since it now does more than simply update display
	def updateDisplay(self):
		self.rawServer.add_task(self.updateDisplay, self.config['display_interval'])
		self.display.status()
		self.checkTimeout()
		self.checkShareRatios()

	def checkTimeout(self):
		# self.torrentInfo only contains torrents that haven't yet finished.
		for torrent, info in self.torrentInfo.items():
			if not info.started or not info.finishTimeout:
				continue
			if info.started + info.finishTimeout > time.time():
				continue
			self.display.error('Torrent did not finish in time', torrent)
			self.shutdownTorrent(torrent)
			self.failed(torrent, True)
		# end for

	def checkShareRatios(self):
		for torrent, info in self.sharingTorrents.items():
			stats = torrent.get_status()
			# TODO: allow share ratio to be overridden
			# TODO: save upTotal in feedStore so that we don't get dinged for restarting.
			if stats['upTotal'] >= stats['downTotal']:
				# TODO: rename finishedTorrent to shutdownTorrent:
				self.shutdownTorrent(torrent)
				if info.handler:
					info.handler.shutdownTorrent(torrent)
				del self.sharingTorrents[torrent]
			# end if
		# end for

	def getConfig(self, pieces, keepAlive):
		config = Preferences(self.config)
		if pieces:
			config['#only-pieces'] = pieces
		if keepAlive is not None:
			config['#piece-picker-next'] = keepAlive.pickPiece
		return config

	def started(self, torrent):
		info = self.torrentInfo.get(torrent)
		# If we already had the file, we will finish before we ever start
		if not info:
			return
		info.started = time.time()
		self.startKeepAlive(info, torrent)
		self.display.startedTorrent(torrent, info)

	def finished(self, torrent):
		info = self.torrentInfo.pop(torrent)
		# If we already had the file, we will finish before we ever start
		if not info.started:
			self.display.startedTorrent(torrent, info)
		def afterFlushed():
			if info.handler and info.handler.finishedTorrent(torrent):
				# We cannot shutdown the torrent immediately, so we schedule it.
				self.rawServer.add_task(self.shutdownTorrent, 0, (torrent,))
				return
			self.sharingTorrents[torrent] = info
		if info.keepAlive:
			info.keepAlive.finish(torrent, afterFlushed)
		else:
			afterFlushed()
		# end if

	def failed(self, torrent, unused_is_external):
		info = self.torrentInfo.pop(torrent)
		if info.handler:
			info.handler.failedTorrent(torrent)
		self.display.failedTorrent(torrent)

	def error(self, *args):
		if len(args) == 2:
			level, text = args
			torrent = None
		else:
			torrent, level, text = args
		getattr(self.display, LEVEL_FUNC[level])(text, torrent)

	def exception(self, torrent, text):
		self.error(torrent, CRITICAL, text)


class LargestFileInfo(object):
	def __init__(self, metainfo):
		super(LargestFileInfo, self).__init__()
		self.metainfo = metainfo
		self.length = metainfo.total_bytes
		self.name = metainfo.name_fs
		# Locate the (piece number, piece offset) for beginning and end of largest file.
		self.begin = (0, 0)
		self.end = divmod(self.length, metainfo.piece_length)
		if not metainfo.is_batch:
			return
		offset = 0
		largest = None
		for i in range(len(metainfo.sizes)):
			size = metainfo.sizes[i]
			if not largest or size > largest[0]:
				largest = (size, i, offset)
			offset += size
		self.length = largest[0]
		self.name = metainfo.files_fs[largest[1]]
		self.begin = divmod(largest[2], metainfo.piece_length)
		self.end = divmod(largest[2] + self.length, metainfo.piece_length)


class BTHashCheck(object):
	def __init__(self, config, metainfo, pathName):
		super(BTHashCheck, self).__init__()
		self.config = config
		self.metainfo = metainfo
		self.pathName = pathName

	def checkPiece(self, piece):
		metainfo = self.metainfo
		if metainfo.is_batch:
			files = [os.path.join(self.pathName, f) for f in metainfo.files_fs]
		else:
			files = [self.pathName]
		filePool = FilePool(self.config['max_files_open'])
		filePool.add_files(files, None)
		storage = Storage(self.config, filePool, zip(files, metainfo.sizes))
		doneFlag = threading.Event()
		config = Preferences(self.config)
		config['check_hashes'] = True
		storageWrapper = StorageWrapper(storage, config, metainfo.hashes, metainfo.piece_length, None, self.statusFunc, doneFlag, None, metainfo.infohash, None, None)
		result = storageWrapper.do_I_have(piece)
		storage.close()
		filePool.close_all()
		return result

	def statusFunc(self, activity=None, fractionDone=0):
		pass
